from typing import Any, Dict

from pydantic import Field

from core.agents.base import BaseAgent
from core.llm.factory import ModelType, create_llm_provider
from core.managers.memory_manager import MemoryUtils
from core.utils.log import mylogger


class MemoryAgent(BaseAgent):

    model_type: ModelType = Field(
        default=ModelType.DOUBAO_FLASH, description="LLM model type"
    )
    short_term_limit: int = Field(default=10, description="Short-term memory limit")
    long_term_limit: int = Field(default=5, description="Long-term memory limit")
    base_dir: str = Field(default=None, description="Memory storage directory")

    def model_post_init(self, __context: Any) -> None:
        """Initialize memory agent."""
        super().model_post_init(__context)
        self.llm = create_llm_provider(model_type=self.model_type)
        self.memory_utils = MemoryUtils(self.name, self.base_dir, self.llm, model_type=self.model_type.value)
        self.memory_utils.SHORT_TERM_LIMIT = self.short_term_limit
        self.memory_utils.LONG_TERM_LIMIT = self.long_term_limit
        mylogger.info(f"MemoryAgent initialized for {self.name} with model {self.model_type.value}")

    async def process(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        operation = input_data.get("operation")

        if operation == "add_short_term":
            content = input_data.get("content", "")
            memory_type = input_data.get("type", "dialogue")
            success = await self.memory_utils.add_short_term_memory(
                content, memory_type
            )
            return {"success": success, "operation": "add_short_term"}

        elif operation == "get_context":
            short_limit = input_data.get("short_limit", 3)
            long_limit = input_data.get("long_limit", 2)
            context = self.memory_utils.get_memory_context(short_limit, long_limit)
            return {"context": context, "operation": "get_context"}

        elif operation == "get_stats":
            stats = self.memory_utils.get_memory_stats()
            return {"stats": stats, "operation": "get_stats"}

        else:
            raise ValueError(f"Unsupported memory operation: {operation}")
